#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import subprocess

# --- 配置区 ---

# 本地数据源的基础目录
LOCAL_BASE_DIR = "/home/atguigu/spark_analysis/data/data_source"
# HDFS上的目标目录
HDFS_TARGET_DIR = "/car_data"
# 需要上传的文件列表
FILES_TO_UPLOAD = ["car_info_data.csv", "car_rank_data.csv"]

# HDFS命令的基础路径，如果 'hdfs' 命令不在系统默认PATH中，请提供完整路径
# 例如: HDFS_CMD = "/opt/module/hadoop-3.1.3/bin/hdfs"
HDFS_CMD = "hdfs"


def execute_command(command, success_msg, error_msg):
    """执行一个shell命令并打印结果 (兼容 Python 3.6)"""
    print(f"[*] 正在执行: {command}")

    # 使用 subprocess.run 执行命令
    # Python 3.6 不支持 capture_output 和 text 参数
    # 我们使用 stdout=subprocess.PIPE 和 stderr=subprocess.PIPE 来替代
    result = subprocess.run(
        command,
        shell=True,
        stdout=subprocess.PIPE,  # 捕获标准输出
        stderr=subprocess.PIPE   # 捕获标准错误
    )

    if result.returncode == 0:
        print(f"[SUCCESS] {success_msg}")
        return True
    else:
        print(f"[ERROR] {error_msg}")
        print(f"  - 返回码: {result.returncode}")
        # result.stderr 是字节流 (bytes), 需要解码 (decode) 为字符串 (string)
        error_details = result.stderr.decode('utf-8').strip()
        print(f"  - 错误信息: {error_details}")
        return False


def main():
    """主执行函数"""
    print("--- 开始执行HDFS文件上传脚本 ---")

    # 1. 确保HDFS目标目录存在
    mkdir_command = f"{HDFS_CMD} dfs -mkdir -p {HDFS_TARGET_DIR}"
    execute_command(
        mkdir_command,
        f"HDFS目录 '{HDFS_TARGET_DIR}' 已确保存在。",
        f"创建HDFS目录 '{HDFS_TARGET_DIR}' 失败。"
    )

    print("\n--- 开始上传文件 ---")
    # 2. 遍历并上传每个文件
    for filename in FILES_TO_UPLOAD:
        local_file_path = os.path.join(LOCAL_BASE_DIR, filename)

        # 2.1 检查本地文件是否存在
        if not os.path.exists(local_file_path):
            print(f"[WARNING] 本地文件不存在，跳过上传: {local_file_path}")
            continue

        # 2.2 执行上传命令
        put_command = f"{HDFS_CMD} dfs -put -f {local_file_path} {HDFS_TARGET_DIR}"
        execute_command(
            put_command,
            f"文件 '{filename}' 已成功上传至 '{HDFS_TARGET_DIR}'。",
            f"上传文件 '{filename}' 失败。"
        )

    print("\n--- 脚本执行完毕 ---")


if __name__ == "__main__":
    main()
